package com.soushangxing.comfyuibackendv2.config

import com.soushangxing.comfyuibackendv2.bean.GlobalQueue
import com.soushangxing.comfyuibackendv2.entity.OrderEntity
import com.soushangxing.comfyuibackendv2.entity.OrderErrorEntity
import com.soushangxing.comfyuibackendv2.repository.OrderErrorRepository
import com.soushangxing.comfyuibackendv2.repository.OrderRepository
import com.soushangxing.comfyuibackendv2.repository.ServerRepository
import com.soushangxing.comfyuibackendv2.service.ComfyUIService
import com.soushangxing.comfyuibackendv2.service.GlobalQueueService
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.context.annotation.Configuration
import org.springframework.web.socket.*
import org.springframework.web.socket.client.WebSocketConnectionManager
import org.springframework.web.socket.client.standard.StandardWebSocketClient
import org.springframework.web.socket.config.annotation.EnableWebSocket
import org.springframework.web.socket.handler.TextWebSocketHandler
import java.time.LocalDateTime

@Configuration
@EnableWebSocket
class WebSocketManager {

    @Autowired
    lateinit var orderRepository: OrderRepository
    @Autowired
    lateinit var orderErrorRepository: OrderErrorRepository
    @Autowired
    lateinit var globalQueue: GlobalQueue
    @Autowired
    lateinit var queueService: GlobalQueueService
    @Autowired
    lateinit var comfyUIService: ComfyUIService
    @Autowired
    lateinit var serverRepository: ServerRepository

    private val connections = mutableMapOf<String, WebSocketConnectionManager>()

    private val connectionStatuses = mutableMapOf<String, Boolean>()

    // 添加重连相关的属性
    private val maxReconnectAttempts = 5
    private val reconnectDelay = 5000L // 5秒
    private val reconnectAttempts = mutableMapOf<String, Int>()

    fun connect(serverUrl: String, clientId: String) {
        val serverUrlWithoutProtocol = serverUrl.replace("http://","").replace("https://","")
        val url = "ws://$serverUrlWithoutProtocol/ws?clientId=$clientId"
        val connectionManager = createConnectionManager(url, serverUrl)
        connectionManager.start()
        connections[url] = connectionManager
        connectionStatuses[url] = false  // 初始状态设为 false
    }

    fun isConnected(url: String): Boolean {
        return connectionStatuses[url] ?: false
    }

    private fun createConnectionManager(url: String, serverUrl: String): WebSocketConnectionManager {
        val client = StandardWebSocketClient()
        val webSocketHandler = createWebSocketHandler(url,serverUrl)

        return WebSocketConnectionManager(client, webSocketHandler, url)
    }

    private fun createWebSocketHandler(url: String, serverUrl: String): TextWebSocketHandler {
        return object : TextWebSocketHandler() {


            override fun afterConnectionEstablished(session: WebSocketSession) {
                println("Socket连接已建立: ${session.remoteAddress.toString()}")
                connectionStatuses[url] = true
                serverOnOps(serverUrl)
            }

            override fun handleTransportError(session: WebSocketSession, exception: Throwable) {
                println("Socket连接出错: ${session.remoteAddress}, 错误: ${exception.message}")

            }

            override fun afterConnectionClosed(session: WebSocketSession, status: CloseStatus) {
                println("Socket连接关闭: ${session.remoteAddress}, 状态: $status")

                if (status.toString().contains("The decoded text message was too big") == true) {
                    // Log and ignore
                    // 已知：Socket似乎在某些步骤（涉及较大图片）会有 Socket连接关闭: /192.169.3.7:6006, 状态: CloseStatus[code=1009, reason=The decoded text message was too big for the output buffer and the endpoint does not support partial messages]
                    // 此类问题也属于异常错误，目前的处理是立刻重连，暂时没有找到解决方案。
                    println("Received a large message, which will be ignored to maintain connection.")
                    reconnect(url,serverUrl)
                    return
                }
                handleDisconnection(url, session, serverUrl)
            }

            override fun handleMessage(session: WebSocketSession, message: WebSocketMessage<*>) {
                when (message) {
                    is TextMessage -> processMessage(session, message)
                }
            }

            private fun processMessage(session: WebSocketSession,message: TextMessage) {

//                println("接收到消息: ${message.payload}")

                /*
                     匹配开始执行：
                     接收到消息: {"type": "execution_start", "data": {"prompt_id": "fba7edcb-c367-4a30-806c-90dc3cc79d35", "timestamp": 1728616459949}}
                 */
                val patternStart = """("type": "execution_start", "data": \{"prompt_id": ")([0-9a-fA-F\-]+)""" // 匹配指定的格式和prompt_id

                // 使用正则表达式查找匹配项
                val matchResultStart = patternStart.toRegex().find(message.payload)

                // 如果匹配成功，提取prompt_id
                matchResultStart?.let {
                    val promptId = it.groupValues[2] // groupValues[2] 是括号中第二个匹配的内容，即prompt_id的值
//                    println("提取到开始的的prompt_id是: $promptId")

                    val orderEntity = orderRepository.findFirstByPromptId(promptId)
                    if (orderEntity != null) {
                        orderEntity.promptId = promptId
                        orderEntity.status = "解析中"
                        orderRepository.save(orderEntity)
                    }
                }

                /*
                    匹配正在生成图片(K Sampler 开始运行,但是不绝对)
                 */
                val patternProgress = """("type": "progress", "data": \{"value": 1, "max": \d+, "prompt_id": ")([0-9a-fA-F\-]+)""" // 匹配指定的格式和prompt_id

                // 使用正则表达式查找匹配项
                val matchResultProgress = patternProgress.toRegex().find(message.payload)

                // 提取promptId
                matchResultProgress?.let{
                    val promptId = matchResultProgress.groupValues.get(2)
                    val orderEntity = orderRepository.findFirstByPromptId(promptId)
                    if (orderEntity != null) {
                        orderEntity.promptId = promptId
                        orderEntity.status = "生成中"
                        orderRepository.save(orderEntity)
                    }
                }

                /*
                    匹配出现结果：接收到消息: {"type": "executed", "data": {"node": "9", "display_node": "9", "output": {"images": [{"filename": "ComfyUI_00183_.png", "subfolder": "", "type": "output"}]}, "prompt_id": "fba7edcb-c367-4a30-806c-90dc3cc79d35"}}
                 */
                // 定义正则表达式，用于匹配 filename 和 prompt_id
                val regex = """("filename":\s*"([^"]+)").*?("prompt_id":\s*"([^"]+)")"""

                // 使用正则表达式查找匹配项
                val matchResult = regex.toRegex().find(message.payload)

                // 检查是否匹配成功
                matchResult?.let {
                    // 提取 filename 和 prompt_id 的值
                    val filename = it.groupValues[2]
                    val promptId = it.groupValues[4]
//                    println("提取到的filename是: $filename")
//                    println("提取到的prompt_id是: $promptId")

                    val orderEntity = orderRepository.findFirstByPromptId(promptId)
                    if (orderEntity != null) {
                        println("TASK ${promptId} FINISHED: Result filename: ${filename}")
                        comfyUIService.viewImage(filename,orderEntity.taskId)
                        // 任务判定为完成，从全局队列中将它移除掉，并且自动安排新的任务（如可用）
//                    globalQueue.globalBackendQueue.removeIf { it.promptId == promptId }
                        globalQueue.globalBackendQueue.removeIf { it.taskId == orderEntity.taskId }
                    }

                    queueService.checkQueue()
                }

                /*
                    匹配出现错误{"type": "execution_error", "data": {"prompt_id": "0c6e72db-f622-449f-b9ea-7e52167bd571",
                 */
                val patternError = """("type": "execution_error", "data": \{"prompt_id": ")([0-9a-fA-F\-]+)""" // 匹配指定的格式和prompt_id

                // 使用正则表达式查找匹配项
                val matchResultError = patternError.toRegex().find(message.payload)

                // 如果匹配成功，提取prompt_id
                matchResultError?.let {
                    val promptId = it.groupValues[2] // groupValues[2] 是括号中第二个匹配的内容，即prompt_id的值
                    println("TASK $promptId HAS AN ERROR: "+message.payload)

                    val orderEntity = orderRepository.findFirstByPromptId(promptId)?:OrderEntity()
                    orderEntity.promptId = promptId
                    orderEntity.status = "任务执行错误"
                    orderRepository.save(orderEntity)

                    val orderErrorEntity = OrderErrorEntity()
                    orderErrorEntity.promptId = orderEntity.promptId
                    orderErrorEntity.taskId = orderEntity.taskId
                    orderErrorEntity.userId = orderEntity.userId
                    orderErrorEntity.createTime = LocalDateTime.now()
                    orderErrorEntity.message = message.payload
                    orderErrorRepository.save(orderErrorEntity)

                    // 任务判定为出错，从全局队列中将它移除掉，并且自动安排新的任务（如可用）
                    globalQueue.globalBackendQueue.removeIf { it.promptId == promptId }
                    queueService.checkQueue()

                }
            }
        }
    }


    private fun handleDisconnection(url: String, session: WebSocketSession, serverUrl: String) {
        connectionStatuses[url] = false

//        清理相关资源
        connections.remove(url)
        reconnectAttempts.remove(url)
        // 处理队列、数据库中服务器数据
        serverOffOps(serverUrl,"故障")

    }

    private fun reconnect(url: String, serverUrl: String) {
        try {
            connections[url]?.stop()
            val connectionManager = createConnectionManager(url, serverUrl)
            connectionManager.start()
            connections[url] = connectionManager
        } catch (e: Exception) {
            println("重连失败: ${e.message}")
        }
    }

    fun disconnect(serverUrl: String) {
        try {
            connections[serverUrl]?.stop()
        } catch (e: Exception) {
            println("断开连接时发生错误: ${e.message}")
        } finally {
            connections.remove(serverUrl)
            connectionStatuses[serverUrl] = false
            reconnectAttempts.remove(serverUrl)

            serverOffOps(serverUrl,"停止")
        }
    }

    fun serverOnOps(serverUrl: String){
        val server = serverRepository.findFirstByUrl(serverUrl)
        if (server != null) {
            server.status = "运行中"
            serverRepository.save(server)

            // 处理全局 globalServer 数组
            val existingServerIndex = globalQueue.globalServer.indexOfFirst { it.url == server.url }
            if (existingServerIndex == -1) {
                // 如果不存在相同 URL 的服务器，则添加新服务器
                globalQueue.globalServer.add(server)
            } else {
                // 如果存在相同 URL 的服务器，则更新其状态
                globalQueue.globalServer[existingServerIndex].status = "运行中"
            }
        }
    }
    fun serverOffOps(serverUrl: String, status:String){
        val server = serverRepository.findFirstByUrl(serverUrl)
        if (server != null) {
            server.status = status
            serverRepository.save(server)

            // 处理全局 globalServer 数组
            val existingServerIndex = globalQueue.globalServer.indexOfFirst { it.url == server.url }
            if (existingServerIndex == -1) {
                // 如果不存在相同 URL 的服务器，则添加新服务器
                globalQueue.globalServer.add(server)
            } else {
                // 如果存在相同 URL 的服务器，则更新其状态
                globalQueue.globalServer[existingServerIndex].status = status
            }
        }
    }
}
